﻿module AssemblyCode

open TinyJava.TAC
open Environment
open System.Collections

type IdType =
| Temp
| BlockVar
| Field

type AssemblyArithOp =
| Add
| Sub
| Mul
| Div

type CmpState =
| Greater
| Geq
| Less
| Leq
| Eq
| Neq

type AssemblyGenerator =
  {
    GeneratedCode         : (string * string) list
    CurrentClassTable     : ClassTable
    CurrentMethodTable    : MethodInfo
    CurrentMethodTAC      : TACGenerator
    SymbolTable           : Globals
    CmpState              : CmpState
  }

let (-) (s1 : string) (s2 : string) : string = s1.Replace(s2,"")

let add_statement (stm : string * string) (g : AssemblyGenerator) : AssemblyGenerator =
  {g with GeneratedCode = g.GeneratedCode @ [stm]}

let set_cmp_state (state : CmpState) (g : AssemblyGenerator) : AssemblyGenerator =
  {g with CmpState = state}

let add_statements (stms : (string * string) list) (g : AssemblyGenerator) : AssemblyGenerator =
  match stms with
  | [] -> g
  | stm :: stms -> 
      let g = add_statement stm g
      add_statements stms g

let is_block_var (id : Id) (m_table : MethodInfo) : bool =
  let block_table = m_table.Body
  block_table.Value.Value.Vars |> List.exists(fun (n,_) -> id = n)

let id_type (id : Id) (g : AssemblyGenerator) : IdType =
  if id.StartsWith(temp_prefix) then
    Temp
  elif is_block_var id (g.CurrentMethodTable) then
    BlockVar
  else
    Field

let find_temp_offset (t : Id) : int =
  let string_index = t - temp_prefix
  let t_idx = int string_index
  -4 * (t_idx + 1)

let find_blockvar_offset (v : Id) (g : AssemblyGenerator) : int =
  let v_idx = g.CurrentMethodTable.Body.Value.Value.Vars |> List.findIndex(fun (id,_) -> v = id)
  let temp_n = g.CurrentMethodTAC.TempList.Length
  -4 * (temp_n + 1 + v_idx)

let find_field_offset (f : Id) (g : AssemblyGenerator) : int =
  let f_idx = g.CurrentClassTable.Fields |> Seq.findIndex(fun ivs -> ivs.Key = f)
  8 * (f_idx + 1)
    
let reg_offset (reg : string) (offset : int) : string =
  if offset = 0 then
    "[" + reg + "]"
  elif offset > 0 then 
    "[" + reg + " + " + (string (abs offset)) + "]"
  else 
    "[" + reg + " - " + (string (abs offset)) + "]"

let comment (c : string) : string =
  "# " + c    

let mov (arg1 : string) (arg2 : string) : string * string = 
  "mov " + arg1 + ", " + arg2,comment (arg1 + " = " + arg2)

let add (arg1 : string) (arg2 : string) : string * string =
  "add " + arg + ", " + arg2,comment (arg1 + " = " + arg1 + " + " + arg2)

let imul (arg1 : string) (arg2 : string) : string * string =
  "imul " + arg + ", " + arg2,comment (arg1 + " = " + arg1 + " * " + arg2)

let sub (arg1 : string) (arg2 : string) : string * string =
  "sub " + arg + ", " + arg2,comment (arg1 + " = " + arg1 + " - " + arg2)

let idiv (arg : string) : string * string =
  "idiv " + arg ,comment (arg1 + " = " + arg1 + " / " + arg2)

let modulus (arg1 : string) (arg2 : string) : string * string =
  "mod " + arg + ", " + arg2,comment (arg1 + " = " + arg1 + " % " + arg2)

let push (arg : string) : string * string =
  "push " + arg,comment ("push " + arg)

let pop (arg : string) : string * string =
  "pop " + arg,comment ("pop " + arg)

let call (arg : f_name) : string * string =
  "call " + arg,comment ("call " + arg)

let comment_stm (arg : string) : string * string =
  "",comment arg

let cdq() : string * string =
  "cdq ","Sign extend EAX"

let cmp (arg1 : string) (arg2 : string) : string * string =
  "cmp " + arg1 + ", " + arg2,comment ("compare " + arg1 + " with " + arg2)

let neg (arg : string) : string * string =
  "neg " + arg,comment("-" + arg)

let not (arg : string) : string * string =
  "not " + arg,comment("!" + arg)

let call_code (args: string list) (f_name : string) (g : AssemblyGenerator) : AssemblyGenerator =
  let g = add_statement (comment_stm "Caller save block") g
  let g = add_statements [push "eax";push "ecx";push "edx"] g
  let push_args = 
    if args.IsEmpty then
      None
    else
      let push_stms =
        [for arg in args do
            let arg_offset = find_temp_offset arg
            let acces_arg = reg_offset "ebp" arg_offset
            yield push access_arg]
      Some push_stms
  
  let g = add_statement (comment_stm "Caller return block") g
  match push_args with
  | None ->
      let g = add_statement (call f_name) g
      let g = add_statements [pop "eax";pop "ecx";pop "edx"] g
      g
  | Some stms ->
      let g = add_statements push_args g
      let g = add_statement (call f_name) g
      let offset = 4 * args.Length
      let g = add_statements [add "esp" offset;pop "eax";pop "ecx";pop "edx"] g
      g

let comparison_code (g : AssemblyGenerator) : AssemblyGenerator =
  let t0_offset = find_temp_offset t0
  let t1_offset = find_temp offset t1
  let access_t0 = reg_offset "ebp" t0_offset
  let access_t1 = reg_offset "ebp" t1_offset
  let g = add_statement (mov "eax" access_t0)
  let g = add_statement (cmp "eax" access_t1)
  g

let asm_map_statement (g : AssemblyGenerator) (s : TACStm) : AssemblyGenerator =
  match s with
  | MemoryOp(mo) ->
      match mo with
      | Assign(t0,t1) ->
          let t0_id_type = id_type t0 g
          let t1_id_type = id_type t1 g
          match t0_id_type,t1_id_type with
          | Temp,Temp -> //t0 = t1
              let t1_offset = find_temp_offset t1
              let t0_offset = find_temp_offset t0
              let access_t0 = reg_offset "ebp" t0_offset
              let access_t1 = reg_offset "ebp" t1_offset
              let g = add_statement (mov "eax" access_t1) g //eax = t1
              let g = add_statement (mov access_t0 "eax") g //t0 = eax
              g
          | BlockVar,Temp -> //v = t1
              let t0_offset = find_blockvar_offset t0 g
              let t1_offset = find_temp_offset t1
              let access_t0 = reg_offset "ebp" t0_offset
              let access_t1 = reg_offset "ebp" t1_offset
              let g = add_statement (mov "eax" access_t1) g  //eax = t1
              let g = add_statement (mov access_t0 "eax") g  //v = t1
              g
          | Field,Temp -> //this.v = t1
              let this_arg = "[ebp + 8]"
              let t0_offset = find_field_offset t0 g
              let t1_offset = find_temp_offset t1
              let access_t0 = reg_offset "eax" t0_offset
              let access_t1 = reg_offset "ebp" t1_offset
              let g = add_statement (mov "eax" this_arg) g //eax = this
              let g = add_statement (mov access_t0 access_t1) g //this.v = t1
              g

          | Temp,BlockVar -> //t0 = v
              let t1_offset = find_blockvar_offset t1 g
              let t0_offset = find_temp_offset t0
              let access_t0 = reg_offset "ebp" t0_offset
              let access_t1 = reg_offset "ebp" t1_offset
              let g = add_statement (mov "eax" access_t1) g  //eax = v
              let g = add_statement (mov access_t0 "eax") g  //t0 = v
              g
          | Temp,Field -> //t0 = this.v
              let this_arg = "[ebp + 8]"
              let t1_offset = find_field_offset t1 g
              let t0_offset = find_temp_offset t0
              let access_t0 = reg_offset "eax" t0_offset
              let access_t1 = reg_offset "ebp" t1_offset
              let g = add_statement (mov "eax" this_arg) g //eax = this
              let g = add_statement (mov access_t0 access_t1) g //this.v = t0
              g
          | _ -> failiwht "Bug in Assembly Generator: One side of the assignment should be a temporary variable"
        | AssignOp(t,op) ->
            let t_offset = find_temp_offset t
            let access_t = reg_offset "ebp" t_offset
            match op with
            | Add(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let t1_offset = find_temp offset t1
                let access_t0 = reg_offset "ebp" t0_offset
                let access_t1 = reg_offset "ebp" t1_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (add "eax" access_t1) g
                let g = add_statement (mov access_t "eax") g
                g

            | StrCat(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let access_t0 = reg_offset "ebp" t0_offset
                let g = call_code [t0;t1] "__strCat" g
                let g = add_statement (mov access_t "eax")
                g
            | Mul(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let t1_offset = find_temp offset t1
                let access_t0 = reg_offset "ebp" t0_offset
                let access_t1 = reg_offset "ebp" t1_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (imul "eax" access_t1) g
                let g = add_statement (mov access_t "eax") g
                g
            | Sub(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let t1_offset = find_temp offset t1
                let access_t0 = reg_offset "ebp" t0_offset
                let access_t1 = reg_offset "ebp" t1_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (sub "eax" access_t1) g
                let g = add_statement (mov access_t "eax") g
                g
            | Div(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let t1_offset = find_temp offset t1
                let access_t0 = reg_offset "ebp" t0_offset
                let access_t1 = reg_offset "ebp" t1_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement cdq g
                let g = add_statement (mov "ecx" access_t1) g
                let g = add_statement (idiv "ecx") g
                let g = add_statement (mov access_t "eax") g
                g
            | Mod(t0,t1) ->
                let t0_offset = find_temp_offset t0
                let t1_offset = find_temp offset t1
                let access_t0 = reg_offset "ebp" t0_offset
                let access_t1 = reg_offset "ebp" t1_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement cdq g
                let g = add_statement (mov "ecx" access_t1) g
                let g = add_statement (idiv "ecx") g
                let g = add_statement (mov access_t "edx") g
                g
            | Less(t0,t1) ->
                let g = add_statement (comment_stm "t0 < t1")
                let g = comparison_code g
                let g = set_cmp_state Less g
                g

            | Leq(t0,t1) ->
                let g = add_statement (comment_stm "t0 <= t1")
                let g = comparison_code g
                let g = set_cmp_state Leq g
                g
            | Greater(t0,t1) ->
                let g = add_statement (comment_stm "t0 > t1")
                let g = comparison_code g
                let g = set_cmp_state Greater g
                g
            | Geq(t0,t1) ->
                let g = add_statement (comment_stm "t0 >= t1")
                let g = comparison_code g
                let g = set_cmp_state Geq g
                g
            | Eq(t0,t1) ->
                let g = add_statement (comment_stm "t0 == t1")
                let g = comparison_code g
                let g = set_cmp_state Equal g
                g
            | Neq(t0,t1) ->
                let g = add_statement (comment_stm "t0 != t1")
                let g = comparison_code g
                let g = set_cmp_state Neq g
                g
            | Sign(t0) ->
                let t0_offset = find_temp_offset t0
                let access_t0 = reg_offset "ebp" t0_offset
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (neg "eax") g
                g
            | Not(t0) ->
                let t0_offset = find_temp_offset t0
                let access_t0 = reg_offset "ebp" t_offset0
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (not "eax") g
                g
            | Length(t0) ->
                let t0_offset = find_temp_offset t0
                let access_t0 = reg_offset "ebp" t0_offset
                let access_array = reg_offset "eax" -4
                let g = add_statement (mov "eax" access_t0) g
                let g = add_statement (mov "ecx" access_array) g
                let g = add_statement (mov access_t "ecx")
                g
            | NewObject(t) ->
            | NewArray(_,t) ->

  | _ -> g




